import os
from config import Constant
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

plt.rcParams['font.sans-serif'] = ['SimHei']


class PlotUtils:
    @staticmethod
    def plot(truth, prediction, epoch, show=False):
        truth = truth.flatten()
        prediction = prediction.flatten()
        tp = []
        end = truth.shape[0] - Constant.OUTPUT_WINDOW
        tp.extend(truth[:end])
        tp.extend(prediction[end:])
        # x = pd.date_range('20210401', '20220401', freq='MS')
        x = pd.date_range(end='20220301', periods=prediction.shape[0], freq='MS')
        plt.plot(x, prediction.flatten(), color='g', label='prediction')
        plt.plot(x, tp, color='red')
        plt.plot(x[:end], truth[:end], color='blue', label=f'{Constant.DEPTH}cm湿度(kg/m2)')
        plt.title(f'{Constant.DEPTH}cm湿度单步预测')
        plt.ylabel('湿度(kg/m2)')
        plt.xlabel('时间')
        plt.grid(True, which='both')
        plt.legend()
        if show:
            plt.show()
        else:
            save_path = os.path.join('graph', f'epoch_{epoch}.png')
            plt.savefig(save_path)
        plt.close()


if __name__ == "__main__":
    t10 = [14.370000039339066, 14.179999925494194, 14.770000010281802, 17.659999997615813, 22.1, 20.360000149607657,
           21.00999983549118, 16.930000039041044, 13.90999991029501, 13.139999908208848, 12.450000178813935,
           12.0999998909235, 15.870000000000001]
    p10 = [10.458335631489755, 10.82500077843666, 16.898786941468718, 17.92825955659151, 17.014129949808122,
           17.11471435725689, 15.519352773539723, 13.494019918441772, 12.99895611345768, 12.832791931927204,
           13.28411540210247, 12.783569735586644, 13.737919403016567]

    t40 = [50.37999982565642, 49.839999866187576, 47.76999994210899, 47.62999997027219, 61.7, 56.11000013709069,
           59.260000218451026, 54.8399999499321, 52.17000004440546, 52.140000152587895, 52.140000152587895,
           52.140000152587895, 45.705000000000005]
    p40 = [48.11329191610218, 46.10129105592147, 44.329472974389795, 43.186626728549605, 44.50467934515328,
           49.35773971185088, 46.67706062136218, 48.00901423893869, 48.6802445384115, 47.842833600714805,
           49.619423390179875, 51.03447835028172, 48.96141332142055]

    t100 = [82.42999964565038, 81.9399993494153, 79.85999972328544, 76.28000016912819, 85.27000005751849,
            85.63999962866306, 87.32000064432621, 92.52000009030104, 93.43000064045191, 93.45, 93.45, 93.45, 67.905]
    p100 = [77.5350136524439, 76.44329985022544, 75.06973135113716, 73.06065027616918, 69.52602498270572,
            71.01514021933079, 71.0585111182183, 73.93546657048167, 79.62677612990142, 82.04146545678377,
            83.33213596343994, 80.76530533432961, 81.00657815605402]

    t200 = [165.91999997496603, 165.91999997496603, 165.91000000596046, 165.71000001788138, 165.46000000238416,
            165.14999999046324, 164.8499999475479, 164.59000002384184, 164.48999996900557, 164.48, 164.48, 164.48,
            166.51999999999998]
    p200 = [165.95847987174986, 165.9950416922569, 165.99773468375204, 165.96993981361388, 165.8320041143894,
            165.60451664447783, 165.39004678964614, 165.19110746145248, 165.02055121183395, 165.05419940948485,
            165.12987286567687, 165.1002882003784, 165.10544583082196]

    fig, ax = plt.subplots(4, 1, sharex=True)

    # plt.xticks(fontsize=16)
    # plt.yticks(fontsize=16)
    ax = ax.flatten()

    tp = []
    end = 13 - Constant.OUTPUT_WINDOW
    t = t10
    p = p10
    # ax_now = ax[0]
    tp.extend(t[:end])
    tp.extend(p[end:])
    x = pd.date_range(end='20220401', periods=13, freq='MS')
    ax[0].plot(x, p, color='g', label='prediction')
    ax[0].plot(x, tp, color='red')
    ax[0].plot(x[:end], t[:end], color='blue', label=f'10cm湿度(kg/m2)')
    # ax[0].set_title(f'10cm湿度多变量单步预测')
    ax[0].set_ylabel('湿度(kg/m2)')
    # ax[0].set_xlabel('时间')
    ax[0].grid(True, which='both')
    ax[0].legend()

    tp = []
    end = 13 - Constant.OUTPUT_WINDOW
    t = t40
    p = p40
    ax_now = ax[1]
    tp.extend(t[:end])
    tp.extend(p[end:])
    x = pd.date_range(end='20220401', periods=13, freq='MS')
    ax_now.plot(x, p, color='g', label='prediction')
    ax_now.plot(x, tp, color='red')
    ax_now.plot(x[:end], t[:end], color='blue', label=f'40cm湿度(kg/m2)')
    # ax_now.set_title(f'40cm湿度多变量单步预测')
    ax_now.set_ylabel('湿度(kg/m2)')
    # ax_now.set_xlabel('时间')
    ax_now.grid(True, which='both')
    ax_now.legend()

    tp = []
    end = 13 - Constant.OUTPUT_WINDOW
    t = t100
    p = p100
    ax_now = ax[2]
    tp.extend(t[:end])
    tp.extend(p[end:])
    x = pd.date_range(end='20220401', periods=13, freq='MS')
    ax_now.plot(x, p, color='g', label='prediction')
    ax_now.plot(x, tp, color='red')
    ax_now.plot(x[:end], t[:end], color='blue', label=f'100cm湿度(kg/m2)')
    # ax_now.set_title(f'100cm湿度多变量单步预测')
    ax_now.set_ylabel('湿度(kg/m2)')
    # ax_now.set_xlabel('时间')
    ax_now.grid(True, which='both')
    ax_now.legend()

    tp = []
    end = 13 - Constant.OUTPUT_WINDOW
    t = t200
    p = p200
    ax_now = ax[3]
    tp.extend(t[:end])
    tp.extend(p[end:])
    x = pd.date_range(end='20220401', periods=13, freq='MS')
    ax_now.plot(x, p, color='g', label='prediction')
    ax_now.plot(x, tp, color='red')
    ax_now.plot(x[:end], t[:end], color='blue', label=f'200cm湿度(kg/m2)')
    # ax_now.title(f'200cm湿度多变量单步预测')
    ax_now.set_ylabel('湿度(kg/m2)')
    ax_now.set_xlabel('时间')
    ax_now.grid(True, which='both')
    ax_now.legend()

    plt.show()
